package com.langchain4j.core.llm.provider.impl;

import com.langchain4j.core.llm.provider.LLMProvider;
import com.langchain4j.core.llm.model.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.*;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.CompletableFuture;

/**
 * 百度文心一言提供者实现
 */
@Slf4j
@Component("baiduProvider")
public class BaiduProvider implements LLMProvider {
    
    @Value("${langchain4j.baidu.api-key}")
    private String apiKey;
    
    @Value("${langchain4j.baidu.secret-key}")
    private String secretKey;
    
    @Value("${langchain4j.baidu.base-url:https://aip.baidubce.com}")
    private String baseUrl;
    
    @Value("${langchain4j.baidu.timeout:30}")
    private Integer timeout;
    
    @Value("${langchain4j.baidu.model:ernie-bot}")
    private String defaultModel;
    
    @Value("${langchain4j.baidu.embedding-model:text-embedding-v1}")
    private String defaultEmbeddingModel;
    
    private final RestTemplate restTemplate = new RestTemplate();
    private String accessToken;
    private LocalDateTime tokenExpireTime;
    
    @Override
    public CompletableFuture<GenerationResult> generate(String prompt, GenerationConfig config) {
        return CompletableFuture.supplyAsync(() -> {
            try {
                String token = getAccessToken();
                String modelName = config.getModelName() != null ? config.getModelName() : defaultModel;
                
                String url = baseUrl + "/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/" + modelName + "?access_token=" + token;
                
                Map<String, Object> requestBody = new HashMap<>();
                List<Map<String, String>> messages = new ArrayList<>();
                
                // 添加系统消息
                if (config.getSystemPrompt() != null && !config.getSystemPrompt().trim().isEmpty()) {
                    Map<String, String> systemMessage = new HashMap<>();
                    systemMessage.put("role", "system");
                    systemMessage.put("content", config.getSystemPrompt());
                    messages.add(systemMessage);
                }
                
                // 添加用户消息
                Map<String, String> userMessage = new HashMap<>();
                userMessage.put("role", "user");
                userMessage.put("content", prompt);
                messages.add(userMessage);
                
                requestBody.put("messages", messages);
                requestBody.put("temperature", config.getTemperature());
                requestBody.put("top_p", config.getTopP());
                requestBody.put("penalty_score", config.getFrequencyPenalty());
                
                HttpHeaders headers = new HttpHeaders();
                headers.setContentType(MediaType.APPLICATION_JSON);
                
                HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);
                
                ResponseEntity<Map> response = restTemplate.exchange(url, HttpMethod.POST, entity, Map.class);
                
                if (response.getStatusCode() == HttpStatus.OK && response.getBody() != null) {
                    Map<String, Object> body = response.getBody();
                    String content = (String) body.get("result");
                    
                    // 计算token使用情况
                    GenerationResult.TokenUsage tokenUsage = GenerationResult.TokenUsage.builder()
                            .promptTokens(estimateTokens(prompt))
                            .completionTokens(estimateTokens(content))
                            .build();
                    
                    return GenerationResult.success(content, modelName, tokenUsage);
                } else {
                    return GenerationResult.failure("API调用失败", modelName);
                }
                
            } catch (Exception e) {
                log.error("Baidu generation failed", e);
                return GenerationResult.failure(e.getMessage(), config.getModelName() != null ? config.getModelName() : defaultModel);
            }
        });
    }
    
    @Override
    public CompletableFuture<EmbeddingResult> embed(String text) {
        return CompletableFuture.supplyAsync(() -> {
            try {
                String token = getAccessToken();
                String url = baseUrl + "/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/" + defaultEmbeddingModel + "?access_token=" + token;
                
                Map<String, Object> requestBody = new HashMap<>();
                requestBody.put("input", text);
                
                HttpHeaders headers = new HttpHeaders();
                headers.setContentType(MediaType.APPLICATION_JSON);
                
                HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);
                
                ResponseEntity<Map> response = restTemplate.exchange(url, HttpMethod.POST, entity, Map.class);
                
                if (response.getStatusCode() == HttpStatus.OK && response.getBody() != null) {
                    Map<String, Object> body = response.getBody();
                    List<Map<String, Object>> data = (List<Map<String, Object>>) body.get("data");
                    
                    if (data != null && !data.isEmpty()) {
                        List<Double> embedding = (List<Double>) data.get(0).get("embedding");
                        List<Float> floatEmbedding = embedding.stream()
                                .map(Double::floatValue)
                                .toList();
                        
                        Integer tokenCount = estimateTokens(text);
                        return EmbeddingResult.success(floatEmbedding, defaultEmbeddingModel, text, tokenCount);
                    }
                }
                
                return EmbeddingResult.failure("Embedding API调用失败", defaultEmbeddingModel, text);
                
            } catch (Exception e) {
                log.error("Baidu embedding failed", e);
                return EmbeddingResult.failure(e.getMessage(), defaultEmbeddingModel, text);
            }
        });
    }
    
    @Override
    public ModelInfo getModelInfo() {
        return ModelInfo.builder()
                .id("baidu-" + defaultModel)
                .name(defaultModel)
                .providerType("baidu")
                .modelType(ModelInfo.ModelType.CHAT)
                .contextLength(4096)
                .supportedParameters(Arrays.asList("temperature", "top_p", "penalty_score"))
                .capabilities(Arrays.asList(
                        ModelInfo.ModelCapability.TEXT_GENERATION,
                        ModelInfo.ModelCapability.CODE_GENERATION,
                        ModelInfo.ModelCapability.TRANSLATION,
                        ModelInfo.ModelCapability.SUMMARIZATION,
                        ModelInfo.ModelCapability.QUESTION_ANSWERING
                ))
                .pricing(ModelInfo.PricingInfo.builder()
                        .inputPricePer1kTokens(0.012)
                        .outputPricePer1kTokens(0.012)
                        .currency("CNY")
                        .billingUnit("per_1k_tokens")
                        .build())
                .status(ModelInfo.ModelStatus.ACTIVE)
                .build();
    }
    
    @Override
    public CompletableFuture<Boolean> isAvailable() {
        return CompletableFuture.supplyAsync(() -> {
            try {
                String token = getAccessToken();
                return token != null && !token.isEmpty();
            } catch (Exception e) {
                log.error("Baidu availability check failed", e);
                return false;
            }
        });
    }
    
    @Override
    public String getProviderType() {
        return "baidu";
    }
    
    /**
     * 获取访问令牌
     */
    private synchronized String getAccessToken() {
        if (accessToken != null && tokenExpireTime != null && LocalDateTime.now().isBefore(tokenExpireTime)) {
            return accessToken;
        }
        
        try {
            String url = baseUrl + "/oauth/2.0/token?grant_type=client_credentials&client_id=" + apiKey + "&client_secret=" + secretKey;
            
            ResponseEntity<Map> response = restTemplate.getForEntity(url, Map.class);
            
            if (response.getStatusCode() == HttpStatus.OK && response.getBody() != null) {
                Map<String, Object> body = response.getBody();
                accessToken = (String) body.get("access_token");
                Integer expiresIn = (Integer) body.get("expires_in");
                
                // 提前5分钟过期
                tokenExpireTime = LocalDateTime.now().plusSeconds(expiresIn - 300);
                
                return accessToken;
            }
        } catch (Exception e) {
            log.error("Failed to get Baidu access token", e);
        }
        
        return null;
    }
    
    /**
     * 估算token数量
     */
    private Integer estimateTokens(String text) {
        return text.length() / 4;
    }
} 